// 21_multithreaded_server.rs
// Rust 多线程 Web 服务器实现

//! 本教程展示如何实现一个简单的多线程 Web 服务器，基于《Rust 程序设计语言》第 21 章的内容。
//! 该服务器能够并发处理多个请求，避免慢请求阻塞其他请求的处理。
//!
//! 主要内容包括：
//! 1. 单线程服务器的局限性
//! 2. 为每个请求创建线程的方法（不推荐）
//! 3. 自定义线程池实现
//! 4. 使用线程池优化服务器性能

use std::{
    fs,
    io::{BufRead, BufReader, Write},
    net::{TcpListener, TcpStream},
    thread,
    time::Duration,
};
use std::sync::mpsc;
use std::sync::Arc;
use std::sync::Mutex;

// 定义任务类型
type Job = Box<dyn FnOnce() + Send + 'static>;

// 线程池结构体
struct ThreadPool {
    workers: Vec<Worker>,
    sender: Option<mpsc::Sender<Job>>,
}

// 工作线程结构体
struct Worker {
    id: usize,
    thread: Option<thread::JoinHandle<()>>,
}

impl ThreadPool {
    /// 创建指定大小的线程池
    ///
    /// # 参数
    ///
    /// * `size` - 线程池中线程的数量
    ///
    /// # 返回值
    ///
    /// 返回一个新的 ThreadPool 实例
    ///
    /// # 示例
    ///
    /// ```
    /// let pool = ThreadPool::new(4);
    /// ```
    pub fn new(size: usize) -> ThreadPool {
        assert!(size > 0);

        let (sender, receiver) = mpsc::channel();
        let receiver = Arc::new(Mutex::new(receiver));
        let mut workers = Vec::with_capacity(size);

        for id in 0..size {
            workers.push(Worker::new(id, Arc::clone(&receiver)));
        }

        ThreadPool {
            workers,
            sender: Some(sender),
        }
    }

    /// 向线程池提交任务
    ///
    /// # 参数
    ///
    /// * `f` - 要执行的闭包任务
    ///
    /// # 示例
    ///
    /// ```
    /// let pool = ThreadPool::new(4);
    /// pool.execute(|| {
    ///     println!("执行任务");
    /// });
    /// ```
    pub fn execute<F>(&self, f: F)
    where
        F: FnOnce() + Send + 'static,
    {
        let job = Box::new(f);

        if let Some(sender) = &self.sender {
            sender.send(job).unwrap();
        }
    }
}

impl Drop for ThreadPool {
    fn drop(&mut self) {
        // 关闭发送端
        drop(self.sender.take());

        // 等待所有线程完成
        for worker in &mut self.workers {
            println!("关闭工作线程 {}", worker.id);

            if let Some(thread) = worker.thread.take() {
                thread.join().unwrap();
            }
        }
    }
}

impl Worker {
    fn new(id: usize, receiver: Arc<Mutex<mpsc::Receiver<Job>>>) -> Worker {
        let thread = thread::spawn(move || loop {
            let message = receiver.lock().unwrap().recv();

            match message {
                Ok(job) => {
                    println!("工作线程 {} 获取任务", id);
                    job();
                }
                Err(_) => {
                    println!("工作线程 {} 断开连接", id);
                    break;
                }
            }
        });

        Worker {
            id,
            thread: Some(thread),
        }
    }
}

fn main() {
    println!("=== 多线程 Web 服务器 ===");
    
    // 启动服务器
    let listener = TcpListener::bind("127.0.0.1:7878").unwrap();
    let pool = ThreadPool::new(4);

    println!("服务器运行在 http://127.0.0.1:7878");
    println!("支持的路由:");
    println!("  - http://127.0.0.1:7878/        - 首页");
    println!("  - http://127.0.0.1:7878/sleep   - 模拟慢请求");
    println!("按 Ctrl+C 停止服务器\n");

    for stream in listener.incoming() {
        let stream = stream.unwrap();

        pool.execute(|| {
            handle_connection(stream);
        });
    }
}

fn handle_connection(mut stream: TcpStream) {
    let buf_reader = BufReader::new(&stream);
    let request_line = match buf_reader.lines().next() {
        Some(line) => match line {
            Ok(line) => line,
            Err(_) => {
                // 如果读取行失败，发送错误响应
                let response = "HTTP/1.1 400 Bad Request\r\nContent-Length: 0\r\n\r\n";
                stream.write_all(response.as_bytes()).unwrap();
                return;
            }
        },
        None => {
            // 如果没有接收到数据，发送错误响应
            let response = "HTTP/1.1 400 Bad Request\r\nContent-Length: 0\r\n\r\n";
            stream.write_all(response.as_bytes()).unwrap();
            return;
        }
    };

    println!("请求: {}", request_line);

    let (status_line, filename) = match &request_line[..] {
        "GET / HTTP/1.1" => ("HTTP/1.1 200 OK", "hello.html"),
        "GET /sleep HTTP/1.1" => {
            println!("处理慢请求，休眠5秒...");
            thread::sleep(Duration::from_secs(5));
            ("HTTP/1.1 200 OK", "hello.html")
        }
        _ => ("HTTP/1.1 404 NOT FOUND", "404.html"),
    };

    let contents = match fs::read_to_string(filename) {
        Ok(content) => content,
        Err(_) => {
            "<!DOCTYPE html>\n\
            <html>\n\
            <head><title>404 Not Found</title></head>\n\
            <body><h1>404 Not Found</h1><p>Sorry, the page you requested was not found.</p></body>\n\
            </html>".to_string()
        }
    };

    // 移除了重复的 contents 定义

    let length = contents.len();
    let response = format!(
        "{}\r\nContent-Length: {}\r\n\r\n{}",
        status_line, length, contents
    );

    stream.write_all(response.as_bytes()).unwrap();
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_thread_pool_creation() {
        let pool = ThreadPool::new(4);
        assert_eq!(pool.workers.len(), 4);
    }

    #[test]
    #[should_panic(expected = "assertion failed")]
    fn test_thread_pool_zero_size() {
        ThreadPool::new(0);
    }

    #[test]
    fn test_thread_pool_execute() {
        let pool = ThreadPool::new(2);
        let counter = Arc::new(Mutex::new(0));
        let counter_clone = Arc::clone(&counter);

        pool.execute(move || {
            let mut num = counter_clone.lock().unwrap();
            *num += 1;
        });

        // 等待任务完成
        thread::sleep(Duration::from_millis(100));

        let result = *counter.lock().unwrap();
        assert_eq!(result, 1);
    }
}